package com.example.mcp.server.controller;

import com.example.mcp.server.model.McpInitializeRequest;
import com.example.mcp.server.model.McpInitializeResponse;
import com.example.mcp.server.model.McpListToolsResponse;
import com.example.mcp.server.model.McpPingRequest;
import com.example.mcp.server.model.McpToolCallRequest;
import com.example.mcp.server.model.McpToolCallResponse;
import com.example.mcp.server.service.McpService;
import java.util.Map;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.CrossOrigin;
import org.springframework.web.bind.annotation.DeleteMapping;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PathVariable;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestHeader;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

@RestController
@RequestMapping("/mcp")
@CrossOrigin(origins = "*")
public class McpController {

    private final McpService mcpService;

    @Autowired
    public McpController(McpService mcpService) {
        this.mcpService = mcpService;
    }

    @PostMapping("/initialize")
    public ResponseEntity<McpInitializeResponse> initialize(@RequestBody McpInitializeRequest request) {
        try {
            McpInitializeResponse response = mcpService.initialize(request);
            return ResponseEntity.ok(response);
        } catch (Exception e) {
            return ResponseEntity.badRequest().build();
        }
    }

    @GetMapping("/tools")
    public ResponseEntity<McpListToolsResponse> listTools(@RequestHeader("X-Session-ID") String sessionId) {
        try {
            if (!mcpService.validateSession(sessionId)) {
                return ResponseEntity.status(401).build();
            }

            McpListToolsResponse response = mcpService.listTools(sessionId);
            return ResponseEntity.ok(response);
        } catch (Exception e) {
            return ResponseEntity.badRequest().build();
        }
    }

    @PostMapping("/tools/call")
    public ResponseEntity<McpToolCallResponse> callTool(
            @RequestHeader("X-Session-ID") String sessionId,
            @RequestBody McpToolCallRequest request) {
        try {
            if (!mcpService.validateSession(sessionId)) {
                return ResponseEntity.status(401).build();
            }

            McpToolCallResponse response = mcpService.callTool(request);
            return ResponseEntity.ok(response);
        } catch (Exception e) {
            return ResponseEntity.badRequest().build();
        }
    }

    @PostMapping("/ping")
    public ResponseEntity<Map<String, String>> ping(@RequestBody McpPingRequest request) {
        return ResponseEntity.ok(Map.of("status", "pong", "id", request.getId()));
    }

    @PostMapping("/session")
    public ResponseEntity<Map<String, String>> createSession() {
        String sessionId = mcpService.createSession();
        return ResponseEntity.ok(Map.of("sessionId", sessionId));
    }

    @DeleteMapping("/session/{sessionId}")
    public ResponseEntity<Void> closeSession(@PathVariable String sessionId) {
        mcpService.closeSession(sessionId);
        return ResponseEntity.ok().build();
    }

    @GetMapping("/health")
    public ResponseEntity<Map<String, String>> health() {
        return ResponseEntity.ok(Map.of(
                "status", "healthy",
                "service", "Spring AI MCP Server",
                "version", "1.0.0"
        ));
    }
}